package com.bysj.agent.rag;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.poi.ApachePoiDocumentParser;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.extern.slf4j.Slf4j;

import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
@Slf4j
public class DynamicContentRetriever {

    private String urls;
    private ObjectMapper objectMapper = new ObjectMapper();

    public DynamicContentRetriever(String urls) {
        this.urls = urls;
    }


    public ContentRetriever contentRetriever() throws JsonProcessingException {
        if (urls == null || urls.trim().isEmpty()) {
            log.warn("无任何URL，无法检索内容");
            return query -> new ArrayList<>();
        }
        
        try {
            // 检查是否是错误响应格式
            if (urls.startsWith("{") && urls.contains("code") && urls.contains("msg")) {
                log.error("收到错误响应而非URL列表: {}", urls);
                return query -> new ArrayList<>();
            }
            
            // 解析URLs JSON数组
            List<String> urlList;
            try {
                urlList = objectMapper.readValue(urls, new TypeReference<List<String>>() {});
            } catch (JsonProcessingException e) {
                log.error("解析URL列表失败: {}", urls, e);
                return query -> new ArrayList<>();
            }
            
            List<Document> documents = new ArrayList<>();
            for(String url : urlList) {
                try {
                    log.info("正在加载文档: {}", url);
                    Document document = UrlDocumentLoader.load(url, new ApacheTikaDocumentParser());
                    documents.add(document);
                    log.info("文档加载成功: {}", url);
                } catch (Exception e) {
                    log.error("加载文档失败: {}", url, e);
                }
            }
            
            if (documents.isEmpty()) {
                log.warn("未能成功加载任何文档");
                return query -> new ArrayList<>();
            }

            //使用内存向量存储
            InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

            //使用默认的文档分割器
            EmbeddingStoreIngestor.ingest(documents, embeddingStore);
            return EmbeddingStoreContentRetriever.from(embeddingStore);
        } catch (Exception e) {
            log.error("内容检索器创建失败", e);
            return query -> new ArrayList<>();
        }
    }
}
